
import os
from dotenv import load_dotenv
import mysql.connector
from neo4j import GraphDatabase
from datetime import datetime
import time

# 加载环境变量
load_dotenv()

def sync_to_neo4j():
    print("=" * 50)
    print(f"开始数据同步 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 50)

    # MySQL连接配置

    mysql_config = {
        'host': os.getenv("DB_HOST"),
        'user': os.getenv("DB_USER"),
        'password': os.getenv("DB_PASSWORD"),
        'database': os.getenv("DB_DATABASE"),
    }
    # Neo4j连接配置
    neo4j_uri = os.getenv("NEO4J_URI")
    neo4j_user = os.getenv("NEO4J_USER")
    neo4j_password = os.getenv("NEO4J_PASSWORD")

    # 连接MySQL
    print("\n[步骤1] 连接MySQL数据库...")
    try:
        mysql_conn = mysql.connector.connect(**mysql_config)
        mysql_cursor = mysql_conn.cursor(dictionary=True)
        print(f"✓ MySQL连接成功! 数据库: {mysql_config['database']}")
    except Exception as e:
        print(f"✗ MySQL连接失败: {str(e)}")
        return

    # 连接Neo4j
    print("\n[步骤2] 连接Neo4j数据库...")
    try:
        neo4j_driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))
        with neo4j_driver.session() as test_session:
            test_session.run("RETURN 1")
        print(f"✓ Neo4j连接成功! URI: {neo4j_uri}")
    except Exception as e:
        print(f"✗ Neo4j连接失败: {str(e)}")
        mysql_cursor.close()
        mysql_conn.close()
        return

    try:
        # Step 1: 同步类别数据
        print("\n[步骤3] 同步类别数据...")
        start_time = time.time()

        mysql_cursor.execute("SELECT * FROM income_expenditure_type")
        categories = mysql_cursor.fetchall()
        print(f"  从MySQL读取 {len(categories)} 条类别记录")

        with neo4j_driver.session() as session:
            # 创建类别节点
            category_count = 0
            for cat in categories:
                session.run("""
                    MERGE (c:Category {type_id: $type_id})
                    SET c.type_name = $type_name,
                        c.is_incomes = $is_incomes,
                        c.is_consumption = $is_consumption,
                        c.is_bring_into_assets = $is_bring_into_assets,
                        c.tab = $tab,
                        c.expect_pay = $expect_pay,
                        c.expr_date = $expr_date
                """, cat)
                category_count += 1
                if category_count % 100 == 0:
                    print(f"  已处理 {category_count}/{len(categories)} 个类别")

            print(f"  ✓ 类别节点创建完成: {category_count} 个节点")

            # 创建父类别关系
            parent_result = session.run("""
                MATCH (c:Category)
                WITH DISTINCT c.tab AS parent_name
                MERGE (p:ParentCategory {name: parent_name})
                RETURN count(p) AS parent_count
            """)
            parent_count = parent_result.single()["parent_count"]
            print(f"  ✓ 父类别节点创建完成: {parent_count} 个节点")

            # 创建分类关系
            relation_result = session.run("""
                MATCH (c:Category), (p:ParentCategory {name: c.tab})
                MERGE (c)-[:HAS_PARENT]->(p)
                RETURN count(*) AS relation_count
            """)
            relation_count = relation_result.single()["relation_count"]
            print(f"  ✓ 分类关系创建完成: {relation_count} 个关系")

        elapsed = time.time() - start_time
        print(f"  类别同步完成! 用时: {elapsed:.2f} 秒")

        # Step 2: 同步交易数据
        print("\n[步骤4] 同步交易数据...")
        start_time = time.time()

        mysql_cursor.execute("SELECT * FROM bill_info")
        transactions = mysql_cursor.fetchall()
        print(f"  从MySQL读取 {len(transactions)} 条交易记录")

        with neo4j_driver.session() as session:
            # 创建用户节点和交易节点
            tx_count = 0
            user_set = set()

            for tx in transactions:
                # 创建/更新用户节点
                if tx['user_id'] not in user_set:
                    session.run("MERGE (u:User {user_id: $user_id})",
                                {'user_id': tx['user_id']})
                    user_set.add(tx['user_id'])

                # 创建交易节点
                session.run("""
                    MERGE (t:Transaction {bill_id: $bill_id})
                    SET t.time = $time,
                        t.money = $money,
                        t.is_income = $is_income,
                        t.note = $note
                """, tx)

                # 创建用户-交易关系
                session.run("""
                    MATCH (u:User {user_id: $user_id}), 
                          (t:Transaction {bill_id: $bill_id})
                    MERGE (u)-[:HAS_TRANSACTION]->(t)
                """, tx)

                # 创建交易-类别关系
                session.run("""
                    MATCH (t:Transaction {bill_id: $bill_id}),
                          (c:Category {type_id: $type_id})
                    MERGE (t)-[:BELONGS_TO]->(c)
                """, tx)

                # 如果有收款用户
                if tx['get_user_id']:
                    if tx['get_user_id'] not in user_set:
                        session.run("MERGE (u2:User {user_id: $user_id})",
                                    {'user_id': tx['get_user_id']})
                        user_set.add(tx['get_user_id'])

                    session.run("""
                        MATCH (t:Transaction {bill_id: $bill_id}),
                              (u2:User {user_id: $get_user_id})
                        MERGE (t)-[:RECEIVED_BY]->(u2)
                    """, tx)

                tx_count += 1
                if tx_count % 100 == 0:
                    print(f"  已处理 {tx_count}/{len(transactions)} 笔交易")

            print(f"  ✓ 交易节点创建完成: {tx_count} 个节点")
            print(f"  ✓ 用户节点创建完成: {len(user_set)} 个节点")

        elapsed = time.time() - start_time
        print(f"  交易同步完成! 用时: {elapsed:.2f} 秒")

        # Step 3: 更新聚合数据
        print("\n[步骤5] 更新聚合统计数据...")
        start_time = time.time()

        with neo4j_driver.session() as session:
            # 更新用户-类别消费统计
            spent_result = session.run("""
                MATCH (u:User)-[:HAS_TRANSACTION]->(t:Transaction)-[:BELONGS_TO]->(c:Category)
                WHERE NOT t.is_income
                WITH u, c, sum(t.money) AS total, count(t) AS count
                MERGE (u)-[r:SPENT_IN]->(c)
                SET r.total_amount = total, r.count = count
                RETURN count(r) AS relation_count
            """)
            spent_count = spent_result.single()["relation_count"]
            print(f"  ✓ 消费关系创建完成: {spent_count} 个关系")

            # 更新用户-类别收入统计
            earned_result = session.run("""
                MATCH (u:User)-[:HAS_TRANSACTION]->(t:Transaction)-[:BELONGS_TO]->(c:Category)
                WHERE t.is_income
                WITH u, c, sum(t.money) AS total, count(t) AS count
                MERGE (u)-[r:EARNED_FROM]->(c)
                SET r.total_amount = total, r.count = count
                RETURN count(r) AS relation_count
            """)
            earned_count = earned_result.single()["relation_count"]
            print(f"  ✓ 收入关系创建完成: {earned_count} 个关系")

        elapsed = time.time() - start_time
        print(f"  聚合数据更新完成! 用时: {elapsed:.2f} 秒")

        print("\n" + "=" * 50)
        print(f"数据同步成功完成!")
        print(f"总用时: {time.time() - start_overall:.2f} 秒")
        print(f"同步记录: {len(categories)} 类别, {len(transactions)} 交易")
        print("=" * 50)

    except Exception as e:
        print("\n" + "!" * 50)
        print(f"同步过程中发生错误: {str(e)}")
        print("!" * 50)
    finally:
        mysql_cursor.close()
        mysql_conn.close()
        neo4j_driver.close()


if __name__ == "__main__":
    start_overall = time.time()
    sync_to_neo4j()